package com.codecrane.core.util;

import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * 基于二叉树实现的权重随机工具类
 * Created by crane on 16/7/8.
 */
public class TreeWeightRandom<T> extends WeightRandom<T> {
    private TreeMap<Integer, T> treeMap;
    private int total;

    public TreeWeightRandom(Map<T, Integer> data) {
        super(data);
    }

    @Override
    public void init() {
        TreeMap<Integer, T> treeMap = new TreeMap<Integer, T>();
        int total = 0;
        for (Map.Entry<T, Integer> each : this.data.entrySet()) {
            final T k = each.getKey();
            final Integer v = each.getValue();
            if (null == v || v.equals(0)) {
                continue;
            }
            total += v;
            treeMap.put(total, k);
        }
        this.treeMap = treeMap;
        this.total = total;
    }

    @Override
    public T random() {
        int s = this.ran.nextInt(total);
        return random(s);
    }

    @Override
    public T random(int ran) {
        if (ran < 0) {
            ran = -ran;
        }
        ran = ran % total;
        ran++;
        SortedMap<Integer, T> map = treeMap.tailMap(ran);
        Integer k = map.firstKey();
        if (null != k) {
            return map.get(k);
        }
        return treeMap.get(treeMap.firstKey());
    }
}
